package relmap

import (
	"micro-rulego/internal/util"
	"sync"
)

// A->B(One-to-many)
// A<-B(One-to-One)
var _ util.Relations = &OneToMany{}

type OneToMany struct {
	entityAToB map[uint32]map[uint32]struct{}
	entityBToA map[uint32]uint32 // 改为每个 entityB 只对应一个 entityA
	lock       sync.RWMutex
}

func NewOneToMany() *OneToMany {
	return &OneToMany{
		entityAToB: make(map[uint32]map[uint32]struct{}),
		entityBToA: make(map[uint32]uint32),
	}
}

func (m *OneToMany) AddRelation(entityA, entityB uint32) bool {
	m.lock.Lock()
	defer m.lock.Unlock()

	// 检查 entityB 是否已经与其他 entityA 关联
	if existingA, exists := m.entityBToA[entityB]; exists {
		if existingA != entityA { // 如果 entityB 已经与其他 entityA 关联，则返回 false
			return false
		}
	}

	// 如果 entityA 没有与 entityB 的关系，则添加
	if _, exists := m.entityAToB[entityA]; !exists {
		m.entityAToB[entityA] = make(map[uint32]struct{})
	}
	m.entityAToB[entityA][entityB] = struct{}{}

	// 建立 entityB 到 entityA 的单向关系
	m.entityBToA[entityB] = entityA

	return true
}

func (m *OneToMany) RemoveRelation(entityA, entityB uint32) {
	m.lock.Lock()
	defer m.lock.Unlock()

	if relatedBs, exists := m.entityAToB[entityA]; exists {
		delete(relatedBs, entityB)
		if len(relatedBs) == 0 {
			delete(m.entityAToB, entityA)
		}
	}

	// 移除 entityB 到 entityA 的关系
	if existingA, exists := m.entityBToA[entityB]; exists && existingA == entityA {
		delete(m.entityBToA, entityB)
	}
}

func (m *OneToMany) GetRelatedToA(entityA uint32) []uint32 {
	m.lock.RLock()
	defer m.lock.RUnlock()

	relatedBs := []uint32{}
	if relations, exists := m.entityAToB[entityA]; exists {
		for entityB := range relations {
			relatedBs = append(relatedBs, entityB)
		}
	}
	return relatedBs
}

func (m *OneToMany) GetRelatedToB(entityB uint32) []uint32 {
	m.lock.RLock()
	defer m.lock.RUnlock()

	entityA, exists := m.entityBToA[entityB]
	if exists {
		return []uint32{entityA}
	}
	return []uint32{}
}
